import os
import sys
from datetime import datetime
import matplotlib
from matplotlib import pyplot as PLT
from scipy.optimize import curve_fit
from scipy.stats import linregress
from datetime import datetime

def Collect(run):
    torque = run.Torque()
    position = run.Position()
    maxTorque = max(torque)
    
    pos_torque = []
    for i, x in enumerate(position):
        if x > 300 and x < 330:
            pos_torque.append(torque[i])
    max_at_B2 = (max(pos_torque))
    
    pos_torque2 = []
    for i, x in enumerate(position):
        if x > 350 and x < 375:
            pos_torque2.append(torque[i])
    max_at_B3 = (max(pos_torque2))
    
    return maxTorque, max_at_B2, max_at_B3

def Linearfnc(x, m, b):
    """ Linear function for fit """
    return m*x + b
    
def LinearFit(date, maxTorque):
    """ Perform linear fit to max torque vs date data 
        Returns fit parameters """
    fit = linregress(date, maxTorque)
    return fit


def FitArray(fit, date):
    """ Produce data array with linear fit equation parameters """
    fitArray = []
    for day in date:
        fitArray.append(Linearfnc(day, fit[0], fit[1]))
    return fitArray


def PlotMaxtvsDate(database, today):
    """ Plot mas torque versus date for each run """
    max_torque = []
    max_torque_Before = []
    max_at_B2 = []
    max_at_B2_Before = []
    max_at_B3 = []
    max_at_B3_Before = []
    date = []
    date_Before = []
    intdate = []
    intdate_Before = []
    filelist = []
    filesBefore = []
    files = []
    linfitBefore = []
    linfit = []
    linfit_B2_Before = []
    linfit_B3_Before = []
    linfit_B2 = []
    linfit_B3 = []
    
    for run in database:
        if run.Direction() != 'in': continue
        if run.Source() != 'weakTh-228': continue
        if run.Station() != 'S5': continue
        if run.NewPosition()[0] > 100: continue
        if max(run.NewPosition()) < 350: continue
        #if max(run.NewPosition()) > 400: continue
        
        maxtorque, maxB2, maxB3 = Collect(run)
        print maxtorque
        intd = int(run.Datetime().strftime('%Y%m%d'))
        
        # sort torque values based on date of deployment
        if intd <= 20130222:
            max_at_B2_Before.append(maxB2)
            max_at_B3_Before.append(maxB3)
            max_torque_Before.append(maxtorque)
            date_Before.append(run.datetime)
        else:
            max_at_B2.append(maxB2)
            max_at_B3.append(maxB3)
            max_torque.append(maxtorque)
            date.append(run.datetime)
            
    # convert dates to ints for construction of plots
    for d in date:
        intd = d - min(date)
        intdate.append(intd.days)
    for d in date_Before:
        intd = d - min(date_Before)
        intdate_Before.append(intd.days)
    
    # fit linear trendlines to torque vs date points
    lin = LinearFit(intdate, max_torque)
    lin_B2 = LinearFit(intdate, max_at_B2)
    lin_B3 = LinearFit(intdate, max_at_B3)
    lin_B = LinearFit(intdate_Before, max_torque_Before)
    lin_B2_B = LinearFit(intdate_Before, max_at_B2_Before)
    lin_B3_B = LinearFit(intdate_Before, max_at_B3_Before)
    
    # get y values for the fit
    linfit = FitArray(lin, intdate)
    linfit_B2 = FitArray(lin_B2, intdate)
    linfit_B3 = FitArray(lin_B3, intdate)
    linfit_B = FitArray(lin_B, intdate_Before)
    linfit_B2_B = FitArray(lin_B2_B, intdate_Before)
    linfit_B3_B = FitArray(lin_B3_B, intdate_Before)
    
    # plot
    fig = PLT.figure(figsize=(15,8),dpi=150)
    box = [0.14, 0.14, 0.76, 0.76]
    ax1 = fig.add_axes(box)
    ax1.set_ylabel('Max Torque (in.oz)')
    ax1.set_xlabel('Date')
    PLT.title('Max Torque Versus Date')
    ax1.grid()
    ax1.scatter(date, max_torque, color = 'g', label = 'Data')
    ax1.plot(date, linfit, '-', color = 'g', label = 'Linear fit: y = %.2f * x + %.2f' % (lin[0], lin[1]))
    ax1.scatter(date_Before, max_torque_Before, color = 'g')
    ax1.plot(date_Before, linfit_B, '-', color = 'black', label = 'Linear fit: y = %.2f * x + %.2f' % (lin_B[0], lin_B[1]))
    #ax1.axvline(datetime.strptime("2013-02-28", "%Y-%m-%d"), colors = 'magenta')
    ax1.text(15, 8, "B refers to before the telescope break", size = 'small')
    h, l = ax1.get_legend_handles_labels()
    PLT.legend(h, l, 'lower left')
    fig.savefig('averagePlots/maxt_vs_date_Update' + today + '.png')
    PLT.close(fig)

    fig2 = PLT.figure(figsize=(15,8),dpi=150)
    box = [0.14, 0.14, 0.76, 0.76]
    ax2 = fig2.add_axes(box)
    ax2.set_ylabel('Max Torque (in.oz)')
    ax2.set_xlabel('Date')
    PLT.title('Max Torque Versus Date for the Second Inner Bend')
    ax2.grid()
    ax2.scatter(date, max_at_B2, color = 'b', label = 'Data')
    ax2.plot(date, linfit_B2, '-', color = 'b', label = 'Linear fit: y = %.2f * x + %.2f' % (lin_B2[0], lin_B2[1]))
    ax2.scatter(date_Before, max_at_B2_Before, color = 'b')
    ax2.plot(date_Before, linfit_B2_B, '-', color = 'black', label = 'Linear fit: y = %.2f * x + %.2f' % (lin_B2_B[0], lin_B2_B[1]))
    ax2.annotate("B refers to before the telescope break", xy=(0.9, 0), size = 'small')
    h, l = ax2.get_legend_handles_labels()
    PLT.legend(h, l, 'lower left')
    fig2.savefig('averagePlots/maxt_vs_date_for_B2_Update-' + today + '.png')
    PLT.close(fig2)
    
    fig3 = PLT.figure(figsize=(15,8),dpi=150)
    box = [0.14, 0.14, 0.76, 0.76]
    ax3 = fig3.add_axes(box)
    ax3.set_ylabel('Max Torque (in.oz)')
    ax3.set_xlabel('Date')
    PLT.title('Max Torque Versus Date for the Third Inner Bend')
    ax3.grid()
    ax3.scatter(date, max_at_B3, color = 'r', marker = 'D', label = 'Data')
    ax3.plot(date, linfit_B3, '-', color = 'r', label = 'Linear fit: y = %.2f * x + %.2f' % (lin_B3[0], lin_B3[1]))
    ax3.scatter(date_Before, max_at_B3_Before, color = 'r', marker = 'D')
    ax3.plot(date_Before, linfit_B3_B, '-', color = 'black', label = 'Linear fit: y = %.2f * x + %.2f' % (lin_B3_B[0], lin_B3_B[1]))
    ax3.annotate("B refers to before the telescope break", xy=(0.9, 0), size = 'small')
    h, l = ax3.get_legend_handles_labels()
    PLT.legend(h, l, 'lower left')
    fig3.savefig('averagePlots/maxt_vs_date_for_B3_Update' + today + '.png')
    PLT.close(fig3)
    
    fig4 = PLT.figure(figsize=(15,8),dpi=150)
    box = [0.14, 0.14, 0.76, 0.76]
    ax4 = fig4.add_axes(box)
    ax4.set_ylabel('Max Torque (in.oz)')
    ax4.set_xlabel('Date')
    PLT.title('Max Torque Versus Date for Bends A6 and A7')
    ax4.grid()
    ax4.scatter(date, max_at_B3, color = 'r', marker = 'D', label = 'Bend A7')
    ax4.plot(date, linfit_B3, '-', color = 'r', label = 'Bend A7 linear fit: y = %.2f * x + %.2f' % (lin_B3[0], lin_B3[1]))
    ax4.scatter(date_Before, max_at_B3_Before, color = 'r', marker = 'D')
    ax4.plot(date_Before, linfit_B3_B, '-', color = 'black', label = 'Bend A7 linear fit: y = %.2f * x + %.2f' % (lin_B3_B[0], lin_B3_B[1]))
    ax4.scatter(date, max_at_B2, color = 'b', label = 'Bend A6')
    ax4.plot(date, linfit_B2, '-', color = 'b', label = 'Bend A6 linear fit: y = %.2f * x + %.2f' % (lin_B2[0], lin_B2[1]))
    ax4.scatter(date_Before, max_at_B2_Before, color = 'b')
    ax4.plot(date_Before, linfit_B2_B, '-', color = 'g', label = 'Bend A6 linear fit: y = %.2f * x + %.2f' % (lin_B2_B[0], lin_B2_B[1]))
    #ax4.axvline(datetime.strptime("2013-02-28", "%Y-%m-%d"), colors = 'magenta')
    PLT.text(0.5, 0.5, "B refers to before the telescope break", size = 'small')
    h, l = ax4.get_legend_handles_labels()
    PLT.legend(h, l, 'lower left')
    fig4.savefig('averagePlots/maxt_vs_date_for_peaks_Update-' + today + '.png')
    PLT.close(fig4)

